package com.frame.geteway.handler;

import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.BlockRequestHandler;
import com.alibaba.csp.sentinel.transport.config.TransportConfig;
import com.frame.common.tools.result.ResultCode;
import com.frame.common.tools.result.ServiceResult;
import org.springframework.cloud.commons.util.InetUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.ServerResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;

/**
 * 自定义限流返回
 *
 * @auther zhangwj
 * @date 2021/4/2 下午2:42
 */
@Component
public class CustomBlockRequestHandler implements BlockRequestHandler {

    @Resource
    private InetUtils inetUtils;

    @PostConstruct
    public void doInit() {
        System.setProperty(TransportConfig.HEARTBEAT_CLIENT_IP, inetUtils.findFirstNonLoopbackAddress().getHostAddress());
    }

    @Override
    public Mono<ServerResponse> handleRequest(ServerWebExchange serverWebExchange,
                                              Throwable throwable) {

        return ServerResponse.status(HttpStatus.TOO_MANY_REQUESTS).contentType(MediaType.APPLICATION_JSON_UTF8)
                .body(BodyInserters.fromObject(ServiceResult.getFailureResult(ResultCode.LIMIT_ACCESS)));
    }
}
